"""
RAG服务性能比较评估工具
用于比较不同RAG实现的性能
"""
import time
import asyncio
import argparse
import json
from pathlib import Path
from typing import Dict, Any, List
from loguru import logger

# 添加项目根目录到Python路径
import sys
import os
sys.path.append(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))

from app.models.schemas import QueryRequest
from app.services.faiss_rag_service import faiss_rag_service
from app.database.connection import db_manager


class RAGComparison:
    """RAG实现比较工具"""
    
    def __init__(self):
        self.db = None
    
    def init_database(self):
        """初始化数据库连接"""
        try:
            db_manager.init_database()
            self.db = db_manager.get_db_session()
            logger.info("数据库连接初始化成功")
        except Exception as e:
            logger.error(f"数据库连接失败: {e}")
            raise
    
    async def compare_implementations(self, question: str) -> Dict[str, Any]:
        """比较不同RAG实现的性能"""
        request = QueryRequest(
            question=question,
            user_id="comparison_test"
        )
        
        results = {}
        
        # 测试FAISS RAG服务
        try:
            start_time = time.time()
            faiss_result = await faiss_rag_service.query(request, self.db)
            faiss_time = time.time() - start_time
            
            results["faiss_rag"] = {
                "response_time": faiss_time,
                "confidence": faiss_result.confidence_score,
                "retrieved_docs": len(faiss_result.retrieved_documents),
                "answer_length": len(faiss_result.answer),
                "answer_preview": faiss_result.answer[:200] + "..." if len(faiss_result.answer) > 200 else faiss_result.answer
            }
        except Exception as e:
            results["faiss_rag"] = {"error": str(e)}
        
        # 测试Milvus RAG服务
        try:
            from app.services.milvus_rag_service import milvus_rag_service
            if milvus_rag_service:
                start_time = time.time()
                milvus_result = await milvus_rag_service.query(request, self.db)
                milvus_time = time.time() - start_time

                results["milvus_rag"] = {
                    "response_time": milvus_time,
                    "confidence": milvus_result.confidence_score,
                    "retrieved_docs": len(milvus_result.retrieved_documents),
                    "answer_length": len(milvus_result.answer),
                    "answer_preview": milvus_result.answer[:200] + "..." if len(milvus_result.answer) > 200 else milvus_result.answer
                }
            else:
                results["milvus_rag"] = {"error": "Milvus service not available"}
        except Exception as e:
            results["milvus_rag"] = {"error": str(e)}
        
        return results
    
    async def batch_compare(self, questions: List[str]) -> Dict[str, Any]:
        """批量比较多个问题"""
        all_results = {}
        summary = {
            "total_questions": len(questions),
            "faiss_rag": {"success": 0, "errors": 0, "avg_time": 0, "avg_confidence": 0},
            "milvus_rag": {"success": 0, "errors": 0, "avg_time": 0, "avg_confidence": 0}
        }
        
        for i, question in enumerate(questions):
            logger.info(f"处理问题 {i+1}/{len(questions)}: {question[:50]}...")
            
            result = await self.compare_implementations(question)
            all_results[f"question_{i+1}"] = {
                "question": question,
                "results": result
            }
            
            # 更新统计信息
            for service_name in ["faiss_rag", "milvus_rag"]:
                if service_name in result:
                    if "error" in result[service_name]:
                        summary[service_name]["errors"] += 1
                    else:
                        summary[service_name]["success"] += 1
                        summary[service_name]["avg_time"] += result[service_name]["response_time"]
                        summary[service_name]["avg_confidence"] += result[service_name]["confidence"]
        
        # 计算平均值
        for service_name in ["faiss_rag", "milvus_rag"]:
            success_count = summary[service_name]["success"]
            if success_count > 0:
                summary[service_name]["avg_time"] /= success_count
                summary[service_name]["avg_confidence"] /= success_count
        
        return {
            "summary": summary,
            "detailed_results": all_results
        }
    
    def save_results(self, results: Dict[str, Any], output_path: str):
        """保存比较结果"""
        Path(output_path).parent.mkdir(parents=True, exist_ok=True)
        with open(output_path, 'w', encoding='utf-8') as f:
            json.dump(results, f, ensure_ascii=False, indent=2)
        logger.info(f"结果已保存到: {output_path}")
    
    def print_summary(self, results: Dict[str, Any]):
        """打印比较结果摘要"""
        if "summary" in results:
            summary = results["summary"]
            print("\n" + "="*60)
            print("RAG服务性能比较结果")
            print("="*60)
            
            for service_name in ["faiss_rag", "milvus_rag"]:
                stats = summary[service_name]
                print(f"\n{service_name.upper()}:")
                print(f"  成功: {stats['success']}")
                print(f"  失败: {stats['errors']}")
                if stats['success'] > 0:
                    print(f"  平均响应时间: {stats['avg_time']:.3f}s")
                    print(f"  平均置信度: {stats['avg_confidence']:.3f}")
        else:
            # 单个问题的结果
            print("\n" + "="*60)
            print("RAG服务比较结果")
            print("="*60)
            
            for service_name, result in results.items():
                print(f"\n{service_name.upper()}:")
                if "error" in result:
                    print(f"  错误: {result['error']}")
                else:
                    print(f"  响应时间: {result['response_time']:.3f}s")
                    print(f"  置信度: {result['confidence']:.3f}")
                    print(f"  检索文档数: {result['retrieved_docs']}")
                    print(f"  答案长度: {result['answer_length']}")
                    print(f"  答案预览: {result.get('answer_preview', 'N/A')}")
    
    def close(self):
        """关闭数据库连接"""
        if self.db:
            self.db.close()


async def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="RAG服务性能比较工具")
    parser.add_argument("--question", type=str, help="单个问题测试")
    parser.add_argument("--questions-file", type=str, help="问题列表文件路径")
    parser.add_argument("--output", type=str, default="evaluation/results/rag_comparison.json", 
                       help="输出文件路径")
    
    args = parser.parse_args()
    
    if not args.question and not args.questions_file:
        print("请提供 --question 或 --questions-file 参数")
        return 1
    
    comparator = RAGComparison()
    comparator.init_database()
    
    try:
        if args.question:
            # 单个问题测试
            print(f"🔍 测试问题: {args.question}")
            results = await comparator.compare_implementations(args.question)
            comparator.print_summary(results)
            comparator.save_results(results, args.output)
            
        elif args.questions_file:
            # 批量测试
            with open(args.questions_file, 'r', encoding='utf-8') as f:
                questions = [line.strip() for line in f if line.strip()]
            
            print(f"🔍 批量测试 {len(questions)} 个问题")
            results = await comparator.batch_compare(questions)
            comparator.print_summary(results)
            comparator.save_results(results, args.output)
        
        print("✅ 比较完成！")
        
    finally:
        comparator.close()
    
    return 0


if __name__ == "__main__":
    sys.exit(asyncio.run(main()))
